import pandas as pd
import numpy as np
import lightgbm as lgb
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_absolute_error
import matplotlib.pyplot as plt
import seaborn as sns

# 读取数据
df = pd.read_csv('your_data.csv')  # 替换为实际文件路径

# 1. 数据预处理
# 确保日期格式正确（假设数据包含日期列）
df['date'] = pd.to_datetime(df['date'])  # 如果没有日期列，需要根据days_since_launch创建

# 2. 构建训练数据集：筛选有完整14天数据的SPU
complete_spu = df.groupby('spu_code')['days_since_launch'].max()
complete_spu = complete_spu[complete_spu >= 14].index
train_df = df[df['spu_code'].isin(complete_spu)].copy()

# 3. 特征工程
# 添加时间特征
train_df['day_of_week'] = train_df['date'].dt.dayofweek  # 周一=0, 周日=6
train_df['is_weekend'] = train_df['day_of_week'].isin([5, 6]).astype(int)

# 添加初始转化率
train_df['initial_conversion_rate'] = train_df['new_sale_qty'] / train_df['before_new_total']
train_df['initial_conversion_rate'] = train_df['initial_conversion_rate'].replace([np.inf, -np.inf], 0).fillna(0)

# 添加时间交互特征
train_df['time_conversion'] = train_df['days_since_launch'] * train_df['initial_conversion_rate']

# 4. 准备训练数据
# 只使用前14天的数据
train_df = train_df[train_df['days_since_launch'].between(1, 14)]

features = [
    'before_new_total',
    'new_sale_qty',
    'initial_conversion_rate',
    'days_since_launch',
    'day_of_week',
    'is_weekend',
    'time_conversion'
]

X = train_df[features]
y = train_df['sale_qty']

# 5. 训练LightGBM模型
# 划分训练集和验证集
X_train, X_val, y_train, y_val = train_test_split(
    X, y, test_size=0.2, random_state=42
)

# 创建数据集
train_data = lgb.Dataset(X_train, label=y_train)
val_data = lgb.Dataset(X_val, label=y_val, reference=train_data)

# 设置模型参数
params = {
    'objective': 'regression',
    'metric': 'mae',
    'boosting_type': 'gbdt',
    'num_leaves': 31,
    'learning_rate': 0.05,
    'feature_fraction': 0.9,
    'bagging_fraction': 0.8,
    'bagging_freq': 5,
    'verbose': 0,
    'seed': 42
}

# 训练模型
model = lgb.train(
    params,
    train_data,
    num_boost_round=1000,
    valid_sets=[val_data],
    callbacks=[
        lgb.early_stopping(stopping_rounds=50),
        lgb.log_evaluation(period=50)
    ]
)

# 6. 模型评估
val_pred = model.predict(X_val)
mae = mean_absolute_error(y_val, val_pred)
print(f"模型验证集MAE: {mae:.2f}")

# 特征重要性可视化
lgb.plot_importance(model, figsize=(10, 6))
plt.title('Feature Importance')
plt.tight_layout()
plt.savefig('feature_importance.png')
plt.show()

# 7. 准备预测数据（新款数据）
# 假设new_items.csv包含需要预测的新款数据
new_items = pd.read_csv('new_items.csv')  # 替换为实际文件路径

# 每个SPU需要预测未来14天的销量
pred_days = pd.DataFrame({
    'days_since_launch': range(1, 15)
})

# 为每个SPU创建14天的预测数据
pred_data = []
for spu in new_items['spu_code'].unique():
    spu_data = new_items[new_items['spu_code'] == spu].iloc[0]
    spu_pred = pred_days.copy()

    # 添加基础特征
    spu_pred['spu_code'] = spu
    spu_pred['before_new_total'] = spu_data['before_new_total']
    spu_pred['new_sale_qty'] = spu_data['new_sale_qty']

    pred_data.append(spu_pred)

pred_df = pd.concat(pred_data, ignore_index=True)

# 添加日期特征（假设以最近日期为上新日）
latest_date = df['date'].max()
pred_df['date'] = pred_df['days_since_launch'].apply(
    lambda x: latest_date + pd.Timedelta(days=x)
)
pred_df['day_of_week'] = pred_df['date'].dt.dayofweek
pred_df['is_weekend'] = pred_df['day_of_week'].isin([5, 6]).astype(int)

# 添加其他特征
pred_df['initial_conversion_rate'] = pred_df['new_sale_qty'] / pred_df['before_new_total']
pred_df['initial_conversion_rate'] = pred_df['initial_conversion_rate'].replace([np.inf, -np.inf], 0).fillna(0)
pred_df['time_conversion'] = pred_df['days_since_launch'] * pred_df['initial_conversion_rate']

# 8. 进行预测
X_pred = pred_df[features]
pred_df['predicted_daily_sales'] = model.predict(X_pred)

# 计算每个SPU的14天总销量预测
total_sales = pred_df.groupby('spu_code')['predicted_daily_sales'].sum().reset_index()
total_sales.rename(columns={'predicted_daily_sales': 'predicted_14d_total_sales'}, inplace=True)

# 合并每日预测和总销量预测
final_pred = pd.merge(pred_df, total_sales, on='spu_code')

# 9. 保存预测结果
# 每日预测
daily_pred = final_pred[['spu_code', 'date', 'days_since_launch', 'predicted_daily_sales']]
daily_pred.to_csv('daily_predictions.csv', index=False)

# 总销量预测
total_pred = final_pred[['spu_code', 'before_new_total', 'new_sale_qty', 'predicted_14d_total_sales']]
total_pred = total_pred.drop_duplicates()
total_pred.to_csv('total_sales_predictions.csv', index=False)

print("预测完成！结果已保存为 daily_predictions.csv 和 total_sales_predictions.csv")


# 10. 可视化预测结果（示例）
def plot_sales_forecast(spu_code):
    spu_data = final_pred[final_pred['spu_code'] == spu_code]

    plt.figure(figsize=(12, 6))
    sns.lineplot(
        x='days_since_launch',
        y='predicted_daily_sales',
        data=spu_data,
        marker='o'
    )
    plt.title(f'SPU {spu_code} - 14天销量预测')
    plt.xlabel('上新后天数')
    plt.ylabel('预测销量')
    plt.grid(True)
    plt.savefig(f'sales_forecast_{spu_code}.png')
    plt.show()
    print(f"总销量预测: {spu_data['predicted_14d_total_sales'].iloc[0]:.0f}")


# 示例：可视化第一个SPU的预测
if not pred_df.empty:
    plot_sales_forecast(pred_df['spu_code'].iloc[0])